'''
Created on Jun 14, 2010

@author: oabalbin
'''
'''
Created on Oct 16, 2009

@author: oabalbin
'''
import operator
import glob
import os
import pickle
import numpy as np
import scipy as sp
from scipy import stats
from optparse import OptionParser
from datetime import datetime
from collections import deque, defaultdict


from common.classes import myArray, gene_copa
import signatures.db.gene_annotations as ga


def open_files(folderpath,filename):
    '''
    Open the file with the expression for the kinases and other genes.
    '''
    myfiles=[]
    for infile in glob.glob( os.path.join(folderpath, filename) ):
        myfiles.append(infile)
    
    return myfiles

def list_of_names(inputfile):
        
        genlist=deque()
        for line in inputfile:        
            line = line.strip('\n')
            fields = line.split('\t')
            #print fields 
            # To skip headers, star reading samples in column 7 of the file
            if fields[0][0] == '#': 
                continue                             
            if fields[0]=='':
                continue
            
            genlist.append(fields[0])
        
        return genlist


def create_myArray(myfile, median_threshold, normal_samples):
    '''
    It creates a sample object and calculates a z-score for the expression level of each gene
    in the sample. Z-score (expvalues - genemedian)/genesd
    Returns a geneArray and also a samplelist with the expression for each sample
    The input matrix is already lo2 normalized    
    '''
    gene_dictionary, sample_dictionary, gene_interval_dict = defaultdict(), defaultdict(), defaultdict()
    celllines =  deque()
    normal_samples_list=[]
    
    first=True

    for line in myfile:
        if line is None:
            continue 
        if line.startswith('#'):
            continue
        # Handeling headers ....
       
        if line.startswith('gene'):
            celllines = line.strip('\n').split('\t')[2:]
            for j,sp in enumerate(celllines):
                sample_dictionary[sp] = j
                print sp,j
                if sp in normal_samples:
                    normal_samples_list.append(j)
                    
            k=0
            continue
        
        # Reading the file.
        fields = line.strip('\n').split('\t') 
        gene_dictionary[k] = fields[0]
        gene_interval_dict[k] = fields[1]
        
        exp=np.zeros(len(fields[2:]))
        expmasked = np.ma.make_mask_none((len(fields[2:]),))        
        exp = np.array(map(float,fields[2:]))
        
        # shift the values by one. In order to avoid the zero coverage values which will be inf when log is applied.    
        expvalues = np.array([exp]) + 1.0
        expvalues = np.ma.array(expvalues,mask=[expmasked])
        
        if first==True:
            geneRawValues = np.array(expvalues)
            first=False
        else:
            geneRawValues = np.concatenate((geneRawValues,expvalues),axis=0)
        
        k+=1
    
    # Pre-process the matrix data in order to filter out low expressing genes.
    # Calculate the median expression value for each gene.
    gene2keep = filter_low_expressing_genes(geneRawValues,median_threshold)
    geneRawValues = geneRawValues[gene2keep,:]
    gene_dictionary = rebuild_gene_dictionary(gene_dictionary,gene2keep)
    gene_interval_dict = rebuild_gene_dictionary(gene_interval_dict,gene2keep)
    
    # Create an GeneArray object
    thisGeneArray = myArray(geneRawValues, gene_dictionary, sample_dictionary, normal_samples_list, gene_interval_dict)

    return thisGeneArray


def filter_low_expressing_genes(geneRawValues,threshold):
    """
    It filters out genes with very low expression values. 
    genes which max expression value  < threshold for all samples are removed.          
    """
    print geneRawValues.shape[0]
    gene_indicators = np.arange(geneRawValues.shape[0])
    exp_cutoff = np.ma.median(geneRawValues, axis=1)
    #median_cutoff = sp.stats.scoreatpercentile(gene_medians,threshold) 
    #exp_cutoff = np.ma.max(geneRawValues, axis=1)
    gene2keep = gene_indicators[exp_cutoff > threshold]
    print len(gene2keep)
    return gene2keep

def rebuild_gene_dictionary(genedict,genes2keep):
    """
    """
    new_gene_dictionary=defaultdict()
    for i in range(len(genes2keep)):
        new_gene_dictionary[i] = genedict[genes2keep[i]]
    
    return new_gene_dictionary
    

def specificity_sensitivity(normal_samples, outlier_samples, all_samples):
    """
    It calculates the specificity and sentivity of the outlier calling
    """
    # Specificity
    len_outlier_samples, len_normal_samples, len_all_samples =  float(len(outlier_samples)), float(len(normal_samples)), float(len(all_samples))
    false_positives = len(normal_samples.intersection(set(outlier_samples)))
    true_negatives = len_normal_samples - false_positives
    specificity = true_negatives / (true_negatives + false_positives)
    
    #Sensitivity
    true_positives = len_outlier_samples - false_positives
    false_negatives =  len_all_samples - len_normal_samples
    sensitivity = true_positives / (true_positives + false_negatives)
    #Prevalence
    prevalence = len_outlier_samples / (len_all_samples - len_normal_samples)
    
    return specificity, sensitivity, prevalence


def copa_at_different_percentiles(thisGeneArray,percentils_list):
    """
    Run copa a different percentiles
    Return the following dictionaries gene_outliersample_dict, copa_rank_dict, copa_score_dict
    """
    gene_outliersample_dict, copa_rank_dict, copa_score_dict = defaultdict(list), defaultdict(list), defaultdict(list)
    
    for percentil in percentils_list:
        gene_outliersample_dict, copa_rank_dict, copa_score_dict = thisGeneArray.get_the_outliers(percentil, gene_outliersample_dict, copa_rank_dict, copa_score_dict)
    
    return gene_outliersample_dict, copa_rank_dict, copa_score_dict


def summarize_results_copa_percentiles(thisGeneArray, gene_outliersample_dict, copa_rank_dict, copa_score_dict):
    """
    Returns a rank list of final outliers.
    """
    
    indicator=np.array(range(len(percentils_list)))
    gene_copascore_values = np.empty(len(copa_rank_dict))
    
    for i in  range(len(copa_rank_dict)):
        if method == 'median':
            mean_copa_rank, mean_copa_score = np.median(copa_rank_dict[i]), np.median(copa_score_dict[i])
        elif method == 'mean':
            mean_copa_rank, mean_copa_score = np.mean(copa_rank_dict[i]), np.mean(copa_score_dict[i])
        #elif method == 'sum':
            
        genename, mean_copa_ind = thisGeneArray.gendict[i], indicator[copa_score_dict[i]>=mean_copa_score][0]
        outlier_samples = gene_outliersample_dict[i][mean_copa_ind]
        # quantifying how many normal samples the outlier hits
        normal_samples = set(thisGeneArray.normal_samples)
        all_samples = thisGeneArray.sampledict.values()
        # Calculation of specificity, sensitivy, and prevalence
        specificity, sensitivity, prevalence = specificity_sensitivity(normal_samples, outlier_samples, all_samples)
        quality_outliers = [specificity, sensitivity, prevalence]
        
        # Make a temporal array with only the copa score for each gene
        gene_copascore_values[i] = mean_copa_score
        # weigth the mean_copa_score for the number of samples in which that gene have a value greater than mean_copa_score
        mean_copa_score = mean_copa_score*len(outlier_samples)
        # Filter outliers according to the level of variance in the normals
        sdt_normals = thisGeneArray.sdv_normals(i)
        sdt_tumors = thisGeneArray.sdv_tumors_mets(i)
        if sdt_normals > sdt_tumors*sdv_factor:
            continue
        
        # Make a list of outliers  
        mean_gene_copa_score_list.append( gene_copa(i,genename,mean_copa_rank,mean_copa_score, mean_copa_ind, outlier_samples, \
                                                    quality_outliers) )
        
    #sorted list of outliers
    final_list_genes = sorted(mean_gene_copa_score_list, key=operator.attrgetter('mean_copa_score'), reverse=True)
    
    return final_list_genes, gene_copascore_values


def identify_outliers_in_sample(thisGeneArray,final_list_genes, outliers_cutoff):
        
    outliers_in_samples=np.zeros((thisGeneArray.expVal.shape[0],thisGeneArray.expVal.shape[1]))
    outliers_in_samples_copascore = np.zeros((thisGeneArray.expVal.shape[0],thisGeneArray.expVal.shape[1]))    
    list_outliers_report = deque()
    i=0
    
    for gene in final_list_genes: 
        if i > outliers_cutoff:
            break
        
        thisoutlier = [ i, gene.indicator, gene.name, gene.mean_copa_score, np.median(thisGeneArray.expVal[gene.indicator]), 
                       np.max(thisGeneArray.expVal[gene.indicator]), np.min(thisGeneArray.expVal[gene.indicator]), gene.quality[0], gene.quality[1],
                       gene.quality[2] ]
        
        # exclude isoforms        
        list_outliers_report.append(thisoutlier)
        i+=1
            
        ### Most important thing  of this function    
        # Get the samples in which this gene is an outlier
        samples = gene_outliersample_dict[gene.indicator][gene.mean_copa_ind]
        outliers_in_samples[gene.indicator,samples] = 1
        # Matrix with copa score for the outliers in each sample
        outliers_in_samples_copascore[gene.indicator,samples] = gene_copascore_values[gene.indicator]
    
    return list_outliers_report, outliers_in_samples, outliers_in_samples_copascore 


def print_list_of_outliers(list_outliers_report, outfile_outliers_list):
    """
    Prints the list of outliers
    """       
    
    for thisoutlier in list_outliers_report:
        print ",".join(map(str,thisoutlier)).replace(',','\t')
        outfile_outliers_list.write(",".join(map(str,thisoutlier)).replace(',','\t')+'\n')
        

def print_matrix_of_outliers(list_of_samples, outliers_in_samples_copascore, outfile_outliers_matrix):
    """
    Prints the matrix of outliers
    """       
    outfile_outliers_matrix.write("geneName"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    for thisoutlier in list_outliers_report:
        eline = list(outliers_in_samples_copascore[thisoutlier[1],:])
        outfile_outliers_matrix.write(thisoutlier[2]+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n') 


def print_outliers_by_samples(thisGeneArray, gene_copascore_values, outliers_in_samples, outfile_outliersinsample):
    """
    Prints a text file with the name of the outliers for each sample.
    """
    
    gene_names_array = np.array(thisGeneArray.get_list_of_genes())
 
    for sample,j in thisGeneArray.sampledict.iteritems():
        # gets the raw copa score (not the weighted by num of samples). Rank using the copa score
        gene_expression = gene_copascore_values[outliers_in_samples[:,j]==1]
        # sort outliers genes according to expression
        sorted_outliers = np.ma.argsort(gene_expression)
        # Get gene names    
        oulier_genes_insample = gene_names_array[outliers_in_samples[:,j]==1]
        outliers_gene_expr = [[str(gene), str(exp)] for gene, exp in zip(oulier_genes_insample[sorted_outliers[::-1]], gene_expression[sorted_outliers[::-1]])]
        outliers_gene_expr = sum(outliers_gene_expr,[])
        print j, sample, outliers_gene_expr[:nout2report]
        
        outfile_outliersinsample.write(str(j)+'\t'+sample+'\t'+",".join(outliers_gene_expr[:nout2report]).replace(',','\t')+'\n')



############################

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-d", "--outfolder", dest="outfolder",
                            help="output folder")
    optionparser.add_option("-s", "--normalsFile", dest="normalsFile",
                            help="annotation file with the normal samples")
    optionparser.add_option("-p", action="store_true", dest="dumpfile",
                            help="use a dump file of the expression matrix. It is passed in annotFile")
 



    
    (options, args) = optionparser.parse_args()
    
    
    t = datetime.now()
    tstamp = t.strftime("%Y_%m_%d_%H_%M")
    
    outfiles = ['outfile_outliers','outfile_outliers_names','outfile_outliers_list','outfile_outliers_matrix']
    time_stamp_outfiles=[]
    for name in outfiles: 
        time_stamp_outfiles.append( tstamp+name+'.txt')
        
    thisfile = open(options.annotFile)
    outfile_outliers= open(time_stamp_outfiles[0],'w')
    outfile_outliersinsample= open(time_stamp_outfiles[1],'w')
    outfile_outliers_list = open(time_stamp_outfiles[2],'w')
    outfile_outliers_matrix = open(time_stamp_outfiles[3],'w')
    array_dump_file = options.annotFile + '_dump_file'
    
    ####################
    ##### Parameters
    
    normal_samples_list = list_of_names(open(options.normalsFile))
    
    gene_outliersample_list=[]
    copa_score_genelist=[]
    mean_gene_copa_score_list=[]
    #percentils_list = [50]
    percentils_list = [80,85,90,95,99] 
    method='mean' # median, mean, sum
    outliers_cutoff=200 #100 #200
    nout2report=20
    expression_cutoff=1.5       # minimum median RPKM of expression
    base_expression_cutoff= 50  # In percentil. It is used to report the expression value of a particular gen in a sample. 
                                # The median usually can be used, however because in many cases median=1, It is maybe desireable to use for e.g the75%. 
    median_shift=1.0
    sdv_factor=0.25             #0.25
    not_report_isoforms=False   # Report gene outlier isoforms
    
    ########### Program
    
    if options.dumpfile:
        thisGeneArray = pickle.load(open(thisfile))
    else:    
        thisGeneArray = create_myArray(thisfile, median_shift, normal_samples_list)
        pickle.dump(thisGeneArray,open(array_dump_file,'w'))
        thisGeneArray.copa_norm_across_samples(expression_cutoff)
        thisGeneArray.copa_norm_in_sample()
    #outliers_cutoff = thisGeneArray.expVal.shape[0]
    
    
    ############ Run the outlier analysis for different percentiles
    
    gene_outliersample_dict, copa_rank_dict, copa_score_dict = copa_at_different_percentiles(thisGeneArray,percentils_list)
    
    ############ Summarize the results from the different percentiles
    
    final_list_genes, gene_copascore_values = summarize_results_copa_percentiles(thisGeneArray,
                                                                                  gene_outliersample_dict, copa_rank_dict, 
                                                                                  copa_score_dict, method)
    
    ############# Identify all outliers for each sample
    
    list_outliers_report, outliers_in_samples, outliers_in_samples_copascore = identify_outliers_in_sample(thisGeneArray,final_list_genes, outliers_cutoff)
    
    
    ### Print outlier list and outlier matrix. 
    print_list_of_outliers(list_outliers_report, outfile_outliers_list)
    print_matrix_of_outliers(thisGeneArray.get_list_of_samples(), outliers_in_samples_copascore, outfile_outliers_matrix)
            
    
    ############# Print sample : outliers gene in that sample
    
    print_outliers_by_samples(thisGeneArray, gene_copascore_values, outliers_in_samples, outfile_outliersinsample)
    
        










